package com.aug_leo.datastructure.binarytree;

import java.util.LinkedList;

/**
 * 树的遍历
 */
public class TreeTraversal {

    public static void main(String[] args) {
        /*
                1
               / \
              2   3
             /   / \
            4   5   6
         */
        TreeNode root = new TreeNode(
                new TreeNode(new TreeNode(4), 2, null),
                1,
                new TreeNode(new TreeNode(5), 3, new TreeNode(6))
        );
        preOrder(root);
        System.out.println();

        inOrder(root);
        System.out.println();

        postOrder(root);
        System.out.println();

        traversal(root);
    }

    /**
     * <h3>前序遍历</h3>
     *
     * @param node 节点
     */
    static void preOrder(TreeNode node) {
        if (node == null) {
            return;
        }
        System.out.print(node.val + "\t"); // 值
        preOrder(node.left); // 左
        preOrder(node.right); // 右
    }

    /**
     * <h3>中序遍历</h3>
     *
     * @param node 节点
     */
    static void inOrder(TreeNode node) {
        if (node == null) {
            return;
        }
        inOrder(node.left); // 左
        System.out.print(node.val + "\t"); // 值
        inOrder(node.right); // 右
    }

    /**
     * <h3>后序遍历</h3>
     *
     * @param node 节点
     */
    static void postOrder(TreeNode node) {
        if (node == null) {
            return;
        }
        postOrder(node.left); // 左
        postOrder(node.right); // 右
        System.out.print(node.val + "\t"); // 值
    }

    /**
     * 遍历二叉树
     * 通过迭代方式实现二叉树的前序、中序、后序遍历
     *
     * @param root 二叉树的根节点
     */
    static void traversal(TreeNode root) {
        // 使用链表实现栈，用于迭代遍历二叉树
        LinkedList<TreeNode> stack = new LinkedList<>();

        // 当前正在处理的节点
        TreeNode curr = root; // 代表当前节点
        TreeNode pop = null; // 最近一次弹栈的元素

        // 当栈不为空或者当前节点不为空时，继续遍历
        while (!stack.isEmpty() || curr != null) {
            if (curr != null) {
                stack.push(curr);
                colorPrintln("前: " + curr.val, 31);
                // 处理左子树
                curr = curr.left;
            } else {
                TreeNode peek = stack.peek();
                // 右子树不存在
                if (peek.right == null) {
                    colorPrintln("中: " + peek.val, 36);
                    pop = stack.pop();
                    colorPrintln("后: " + pop.val, 34);
                }
                // 右子树处理完毕
                else if (peek.right == pop) {
                    pop = stack.pop();
                    colorPrintln("后: " + pop.val, 34);
                }
                // 待处理右子树
                else {
                    colorPrintln("中: " + peek.val, 36);
                    curr = peek.right;
                }
            }
        }
    }

    /*
        31 红
        32 黄
        33 橙
        34 蓝
        35 紫
        36 绿
    */
    private static void colorPrintln(String origin, int color) {
        System.out.printf("\033[%dm%s\033[0m%n", color, origin);
    }
}
